from typing import List
from copy   import deepcopy
from AK import AK
from NS import NS
from SR import SR
from MC import MC

# working mode
AES_ECB = 0
AES_CBC = 1
AES_CTR = 2
AES_CFB = 3
AES_OFB = 4


# AES main module
class AES:
    def __init__(self) -> None:
        self.kernel = AES_kernel()
        

    def encrypt(self, pt: str, key: str, iv: str = None, mode = AES_ECB) -> str:
        if len(pt) % 16 != 0: raise Exception("Illegal text length!")
        pts = []
        res = []

        # get textx by block(16bytes)
        while pt:
            pts.append(pt[:16])
            pt = pt[16:]

        # block encrypt
        for p in pts:
            res.append(self.kernel.encrypt(p,key,iv,mode))
        return res

    def decrypt(self, ct: str, key: str, iv: str = None, mode = AES_ECB) -> str:
        cts = []
        res = []
        
        if type(ct) == str:
            if len(ct) % 16 != 0: raise Exception("Illegal text length!")

            # get textx by block(16bytes)
            while ct:
                cts.append(ct[:16])
                ct = ct[16:]
        else:
            cts = ct

        # block encrypt
        for c in cts:
            res.append(self.kernel.decrypt(c,key,iv,mode))
        return res

    def matrixToHex(self,pts: List[List[List[int]]]) -> str:
        text = ''
        for p in pts:
            text += self.kernel.getText(p)
        return text

    def matrixToText(self,pts: List[List[List[int]]]) -> str:
        text = self.matrixToHex(pts)
        res = ''
        while text:
            res += chr(eval('0x'+text[0:2]))
            text = text[2:]
        return res



# AES kernel module: for calculate
class AES_kernel:
    def __init__(self) -> None:
        self.pt = None
        self.ct = None
        self.iv = None

        # encrypt method dict
        self.en = {
            AES_ECB : self.encrypt_ECB,
            AES_CBC : self.encrypt_CBC,
        }

        # decrypt method dict
        self.de = {
            AES_ECB : self.decrypt_ECB,
            AES_CBC : self.decrypt_CBC,
        }

    # encrypt function
    def encrypt(self, pt: str, key: str, iv: str = None, mode = AES_ECB) -> str:
        # alter: str -> matrix
        self.iv = self.format(iv) if type(iv) == str else iv
        key = self.format(key) if type(key) == str else key
        
        
        pt = self.format(pt) if type(pt) == str else pt


        # encrypt by mode 
        if mode in self.en:
            return self.en[mode](pt,key)
        else:
            raise Exception("Illegal working mode!")
        

    # AES encrypy by ECB
    def encrypt_ECB(self, pt: str, key: str) -> str:
        ak = AK(key)

        # Add round key
        nums = self.ir(pt,ak.getRoundKey(0))
        
        for i in range(1,11):
            # Nibble substitution
            nums = NS().encrypt(nums)

            # Shift Row
            nums = SR().encrypt(nums)
            
            # Mix columns
            if i != 10:
                nums = MC().encrypt(nums)
            
            # Add round key
            nums = self.ir(nums,ak.getRoundKey(i))

        return nums


    # TODO AES encrypy by CBC
    def encrypt_CBC(self, pt: str, key: str) -> str:
        raise NotImplementedError('This method has not been implemented yet.')
    
    # decrypt main function
    def decrypt(self, ct: str, key: str, iv: str = None , mode = AES_ECB) -> List[List[int]]:
        # alter: str -> matrix
        self.iv = self.format(iv) if type(iv) == str else iv
        ct = self.format(ct) if type(ct) == str else ct
        key = self.format(key) if type(key) == str else key

        if mode in self.de:
            return self.de[mode](ct,key)
        else:
            raise Exception("Illegal working mode!")

    # AES decrypy by ECB
    def decrypt_ECB(self, ct: str, key: str) -> str:
        ak = AK(key)
        nums = ct 
        for i in range(10,0,-1):
            # Add round key
            nums = self.ir(nums,ak.getRoundKey(i))

            # Inverse Mix columns
            if i != 10:
                nums = MC().decrypt(nums)
            
            # Inverse Shift Row
            nums = SR().decrypt(nums)

            # Inverse Nibble substitution
            nums = NS().decrypt(nums)
        
        # Add round key
        nums = self.ir(nums,ak.getRoundKey(0))
        
        return nums

    # AES decrypy by CBC
    def decrypt_CBC(self, pt: str, key: str) -> str:
        raise NotImplementedError('This method has not been implemented yet.')

    # alter: string to matrix
    def format(self, text: str) -> List[List[int]]:
        # 字符长度必须为16
        if len(text) != 16: raise Exception("Illegal text length!")

        return [[[[ord(text[j]) for j in range(i*4,i*4+4)] for i in range(4)][y][x] for y in range(4)] for x in range(4)]


    # Initial round
    def ir(self, n1: List[List[int]], n2: List[List[int]]) -> List[List[int]]:
        for i in range(4):  
            for j in range(4):
                n1[i][j] ^= n2[i][j]
        return n1

    # print matrix in hex format
    def printNums(self,nums: List[List[int]]) -> None:
        for i in range(len(nums)):
            print([ hex(nums[i][j]) for j in range(len(nums))])

    # get text from matrix
    def getText(self, nums: List[List[int]]) -> str:
        n = deepcopy(nums)
        k = ''
        for j in range(len(n)):
            for i in range(len(n)):
                byte = hex(n[i][j])
                if len(byte) == 3:  byte = byte[0:2] + '0' + byte[2]
                k += byte[2:]
        return k

    # print maxtrix in text format by byte
    def printText(self, nums: List[List[int]]) -> None:
        n = deepcopy(nums)
        k = ''
        for j in range(4):
            for i in range(4):
                k += hex(n[i][j])
                k += ' '
        print(k)
                
    # alter: matrix to str
    def matrixToString(self,nums: List[List[int]]) -> str:
        res = ''
        for j in range(len(nums)):
            for i in range(len(nums)):
                res += chr(nums[i][j])
        return res
